﻿#ifndef __COMMON_CUH__
#define __COMMON_CUH__

#include <cublas_v2.h>
#include <random>
#include <stdio.h>

#define WARP_SIZE 32
#define DIV_UP(a, b) (((a) + (b) - 1) / (b))
#define CFLOAT4(val) (reinterpret_cast<const float4 *>(&(val))[0])
#define FLOAT4(val) (reinterpret_cast<float4 *>(&(val))[0])
#define ALIGN_16 __align__(16)
#define SHARED_ALIEN16 __shared__ ALIGN_16
#define SWIZZLE_FLOAT4(row, col) ((row) ^ (((col) >> 2) & (0x7)))

#define CHECK_CUDA_ERROR(call)                                                                                     \
    {                                                                                                              \
        cudaError_t err = call;                                                                                    \
        if (err != cudaSuccess) {                                                                                  \
            fprintf(stderr, "CUDA error at %s:%d (%s): %s\n", __FILE__, __LINE__, #call, cudaGetErrorString(err)); \
            exit(1);                                                                                               \
        }                                                                                                          \
    }

#define CHECK_CUBLAS_ERROR(call)                                                     \
    {                                                                                \
        cublasStatus_t err = call;                                                   \
        if (err != CUBLAS_STATUS_SUCCESS) {                                          \
            fprintf(stderr, "CUBLAS error at %s:%d: %d\n", __FILE__, __LINE__, err); \
            exit(1);                                                                 \
        }                                                                            \
    }

void random_init(float *data, int size, int seed = -1, float low = -1.0f, float high = 1.0f) {
    if (seed == -1) {
        std::random_device rd;
        seed = rd();
    }
    std::mt19937 gen(seed);
    std::uniform_real_distribution<float> dis(low, high);
    for (int i = 0; i < size; i++)
        data[i] = dis(gen);
}

float max_diff(float *bench, float *test, int size) {
    float max_diff = 0.0f;
    for (int i = 0; i < size; i++) {
        float diff = fabs(bench[i] - test[i]);
        if (diff > max_diff)
            max_diff = diff;
    }
    return max_diff;
}

class Timer {
private:
    cudaEvent_t start, stop;

public:
    Timer() {
        cudaEventCreate(&start);
        cudaEventCreate(&stop);
    }

    ~Timer() {
        cudaEventDestroy(start);
        cudaEventDestroy(stop);
    }

    void start_timer() {
        cudaEventRecord(start);
    }

    void stop_timer() {
        cudaEventRecord(stop);
        cudaEventSynchronize(stop);
    }

    float get_time() {
        float time;
        cudaEventElapsedTime(&time, start, stop);
        return time;
    }
};

void warmup() {
    const int Ms[] = {256, 512, 1024, 2048, 4096, 8192};
    const int Ns[] = {256, 512, 1024, 2048, 4096, 8192};
    const int Ks[] = {1024, 1024, 1024, 1024, 1024, 1024};
    const int n = sizeof(Ms) / sizeof(Ms[0]);
    for (int i = 0; i < n; ++i) {
        const int M = Ms[i], N = Ns[i], K = Ks[i];
        const int size_A = M * K * sizeof(float);
        const int size_B = K * N * sizeof(float);
        const int size_C = M * N * sizeof(float);
        const float alpha = 1.0f;
        const float beta = 0.0f;

        cublasHandle_t handle;
        CHECK_CUBLAS_ERROR(cublasCreate(&handle));

        float *h_a = (float *)malloc(size_A);
        float *h_b = (float *)malloc(size_B);
        float *h_c = (float *)malloc(size_C);
        float *d_a, *d_b, *d_c;
        CHECK_CUDA_ERROR(cudaMalloc(&d_a, size_A));
        CHECK_CUDA_ERROR(cudaMalloc(&d_b, size_B));
        CHECK_CUDA_ERROR(cudaMalloc(&d_c, size_C));

        random_init(h_a, M * K);
        random_init(h_b, K * N);
        memset(h_c, 0, size_C);

        CHECK_CUDA_ERROR(cudaMemcpy(d_a, h_a, size_A, cudaMemcpyHostToDevice));
        CHECK_CUDA_ERROR(cudaMemcpy(d_b, h_b, size_B, cudaMemcpyHostToDevice));
        CHECK_CUDA_ERROR(cudaMemset(d_c, 0, size_C));

        for (int i = 0; i < 10; ++i)
            CHECK_CUBLAS_ERROR(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                                           N, M, K,
                                           &alpha,
                                           d_b, N,
                                           d_a, K,
                                           &beta,
                                           d_c, N));

        cublasDestroy(handle);
        free(h_a);
        free(h_b);
        free(h_c);
        CHECK_CUDA_ERROR(cudaFree(d_a));
        CHECK_CUDA_ERROR(cudaFree(d_b));
        CHECK_CUDA_ERROR(cudaFree(d_c));
    }
    printf("Warmup done\n");
}

#endif // __COMMON_CUH__
